module upwind_mod
	
	use namelist_mod
	use mesh_mod
	use state_mod
	use tracer_mod
	use parallel_mod

	private
	public upwind_init
	public upwind_final
	public calc_tracer_mass_flux_hori_upwind
	public calc_tracer_mass_flux_vert_upwind
	public do_hori_limiter,&
				 do_vert_limiter

	logical, allocatable :: do_hori_limiter(:)
	logical, allocatable :: do_vert_limiter(:)

contains
	
	subroutine upwind_init()
    
    if(.not. allocated(do_hori_limiter)) allocate(do_hori_limiter(3))
		if(.not. allocated(do_vert_limiter)) allocate(do_vert_limiter(3))
		do_hori_limiter(1:3) = .false.
		do_vert_limiter(1:3) = .false.

		do_hori_limiter(1:3) = upwind_hori_fct
		do_vert_limiter(1:3) = upwind_vert_fct
	
	end subroutine upwind_init

	subroutine upwind_final()

		if (allocated(do_hori_limiter)) deallocate(do_hori_limiter)
		if (allocated(do_vert_limiter)) deallocate(do_vert_limiter)

	end subroutine upwind_final
		
	subroutine calc_tracer_mass_flux_hori_upwind(old_tracer, old_state, old_q, q, qmfx, qmfy, dt, do_limiter)

		type(tracer_type), intent(inout) :: old_tracer
		type(state_type), intent(in) :: old_state
		real(r8), intent(in) :: old_q(mesh%full_lon_lb:mesh%full_lon_ub, &
														      mesh%full_lat_lb:mesh%full_lat_ub, &
														      mesh%full_lev_lb:mesh%full_lev_ub)
		real(r8), intent(in) :: q(mesh%full_lon_lb:mesh%full_lon_ub, &
														  mesh%full_lat_lb:mesh%full_lat_ub, &
														  mesh%full_lev_lb:mesh%full_lev_ub)
		real(r8), intent(out) :: qmfx(mesh%half_lon_lb:mesh%half_lon_ub, &
																  mesh%full_lat_lb:mesh%full_lat_ub, &
																  mesh%full_lev_lb:mesh%full_lev_ub)
		real(r8), intent(out) :: qmfy(mesh%full_lon_lb:mesh%full_lon_ub, &
																  mesh%half_lat_lb:mesh%half_lat_ub, &
																  mesh%full_lev_lb:mesh%full_lev_ub)
		real(8), intent(in) :: dt
		logical, intent(in),optional :: do_limiter
		real(r8) pole, q_max, q_min, pin, qin, pout, qout, wgt_lon, wgt_lat
		integer i, j, k

		associate (fx_low => old_tracer%fx_low, &
               fy_low => old_tracer%fy_low, &
							 q_td   => old_tracer%q_td  , &
							 rin    => old_tracer%rin   , &
							 rout   => old_tracer%rout)

		call calc_hori_flux_upwind(old_tracer, old_state, q, dt, '3rd', qmfx, qmfy)

		if (do_limiter) then
			call calc_hori_flux_upwind(old_tracer, old_state, q, dt, '1st', fx_low, fy_low)

			do k = mesh%full_lev_ibeg, mesh%full_lev_iend
				do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
			  	do i = mesh%full_lon_ibeg, mesh%full_lon_iend
			  		q_td(i,j,k) = old_q(i,j,k) - &
						   dt * ((fx_low(i,j,k) - fx_low(i-1,j,k)) * mesh%le_lon(j) + &
						   	     (fy_low(i,j  ,k) * mesh%le_lat(j  ) - &
						   	      fy_low(i,j-1,k) * mesh%le_lat(j-1))) / mesh%area_cell(j)
					end do
				end do
			end do

			!south pole
		  do k = mesh%full_lev_ibeg, mesh%full_lev_iend
			  j = mesh%full_lat_ibeg
			  pole = sum(fy_low(mesh%full_lon_ibeg:mesh%full_lon_iend,j,k)) * mesh%le_lat(j) / mesh%num_full_lon / mesh%area_cell(j)
			  do i = mesh%full_lon_ibeg, mesh%full_lon_iend
			  	q_td(i,j,k) = old_q(i,j,k) - pole * dt
			  end do
			end do
		  ! North pole
		  do k = mesh%full_lev_ibeg, mesh%full_lev_iend
			  j = mesh%full_lat_iend
			  pole = sum(fy_low(mesh%full_lon_ibeg:mesh%full_lon_iend,j-1,k)) * mesh%le_lat(j-1) / mesh%num_full_lon / mesh%area_cell(j)
			  do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					q_td(i,j,k) = old_q(i,j,k) + pole * dt
			  end do
			end do
			call fill_zonal_halo_cell(q_td, all_halo=.true.)

			do k = mesh%full_lev_ibeg, mesh%full_lev_iend
				do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
					do i = mesh%full_lon_ibeg, mesh%full_lon_iend
						q_max = max(maxval(old_q(i-1:i+1,j,k)), maxval(q_td(i-1:i+1,j,k)), &
											  maxval(old_q(i,j-1:j+1,k)), maxval(q_td(i,j-1:j+1,k)))
						q_min = min(minval(old_q(i-1:i+1,j,k)), minval(q_td(i-1:i+1,j,k)), &
											  minval(old_q(i,j-1:j+1,k)), minval(q_td(i,j-1:j+1,k)))		
						pin = ((max(0.0, qmfx(i-1,j,k) - fx_low(i-1,j,k)) - &
						        min(0.0, qmfx(i  ,j,k) - fx_low(i  ,j,k))) * mesh%le_lon(j) + &
									 (max(0.0, qmfy(i,j-1,k) - fy_low(i,j-1,k)) * mesh%le_lat(j-1) - &
									  min(0.0, qmfy(i,j  ,k) - fy_low(i,j  ,k)) * mesh%le_lat(j  ))) / mesh%area_cell(j)
						qin = (q_max - q_td(i,j,k)) / dt
						if (pin > 0.0) then
							rin(i,j,k) = min(1.0_r8, qin / pin)
						else
							rin(i,j,k) = 0.0_r8
						end if
						pout = ((max(0.0, qmfx(i  ,j,k) - fx_low(i  ,j,k)) - &
						         min(0.0, qmfx(i-1,j,k) - fx_low(i-1,j,k))) * mesh%le_lon(j) + & 
									  (max(0.0, qmfy(i,j  ,k) - fy_low(i,j  ,k)) * mesh%le_lat(j  ) - &
										 min(0.0, qmfy(i,j-1,k) - fy_low(i,j-1,k)) * mesh%le_lat(j-1))) / mesh%area_cell(j)
						qout = (q_td(i,j,k) - q_min) / dt
						if (pout > 0.0) then
							rout(i,j,k) = min(1.0_r8, qout / pout)
						else
							rout(i,j,k) = 0.0_r8
						end if
					end do
				end do
			end do
			! South pole
			do k = mesh%full_lev_ibeg, mesh%full_lev_iend
				pole = 0.0
				j = mesh%full_lat_ibeg
				q_max = max(maxval(old_q(mesh%full_lon_ibeg:mesh%full_lon_iend,j+1,k)), &
				 						maxval(q_td (mesh%full_lon_ibeg:mesh%full_lon_iend,j+1,k)), &
				 									 old_q(mesh%full_lon_ibeg,j,k))
				q_min = min(minval(old_q(mesh%full_lon_ibeg:mesh%full_lon_iend,j+1,k)), &
									  minval(q_td (mesh%full_lon_ibeg:mesh%full_lon_iend,j+1,k)), &
													 old_q(mesh%full_lon_ibeg,j,k))
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					pole = pole - min(0.0, qmfy(i,j,k) - fy_low(i,j,k))
				end do
				pin = pole * mesh%le_lat(j) / mesh%num_full_lon / mesh%area_cell(j)
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					qin = (q_max - q_td(i,j,k)) / dt
					if (pin > 0) then
						rin(i,j,k) = min(1.0, qin / pin)
					else
						rin(i,j,k) = 0.0
					end if
				end do

				pole = 0.0
			  do i = mesh%full_lon_ibeg, mesh%full_lon_iend
			  	pole = pole + max(0.0, qmfy(i,j,k) - fy_low(i,j,k))
			  end do
			  pout = pole * mesh%le_lat(j) / mesh%num_full_lon / mesh%area_cell(j)
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					qout = (q_td(i,j,k) - q_min) / dt
					if (pout > 0) then
						rout(i,j,k) = min(1.0, qout / pout)
					else
						rout(i,j,k) = 0.0
					end if
				end do
			end do
			! North pole
			do k = mesh%full_lev_ibeg, mesh%full_lev_iend
				pole = 0.0
				j = mesh%full_lat_iend
				q_max = max(maxval(old_q(mesh%full_lon_ibeg:mesh%full_lon_iend,j-1,k)), &
				 					  maxval(q_td (mesh%full_lon_ibeg:mesh%full_lon_iend,j-1,k)), &
				 									 old_q(mesh%full_lon_ibeg,j,k))
				q_min = min(minval(old_q(mesh%full_lon_ibeg:mesh%full_lon_iend,j-1,k)), &
										minval(q_td (mesh%full_lon_ibeg:mesh%full_lon_iend,j-1,k)), &
													 old_q(mesh%full_lon_ibeg,j,k))
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					pole = pole + max(0.0, qmfy(i,j-1,k) - fy_low(i,j-1,k))
				end do
				pin = pole * mesh%le_lat(j-1) / mesh%num_full_lon / mesh%area_cell(j)
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					qin = (q_max - q_td(i,j,k)) / dt
					if (pin > 0) then
						rin(i,j,k) = min(1.0, qin / pin)
					else
						rin(i,j,k) = 0.0
					end if
				end do

				pole = 0.0
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					pole = pole - min(0.0, qmfy(i,j-1,k) - fy_low(i,j-1,k))
				end do
				pout = pole * mesh%le_lat(j-1) / mesh%num_full_lon / mesh%area_cell(j)
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					qout = (q_td(i,j,k) - q_min) / dt
					if (pout > 0) then
						rout(i,j,k) = min(1.0, qout / pout)
					else
						rout(i,j,k) = 0.0
					end if
				end do
			end do
			call fill_zonal_halo_cell(rin , all_halo=.true.)
			call fill_zonal_halo_cell(rout, all_halo=.true.)

			do k = mesh%full_lev_ibeg, mesh%full_lev_iend
				do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
					do i = mesh%half_lon_ibeg, mesh%half_lon_iend
						if (qmfx(i,j,k) - fx_low(i,j,k) >= 0.0) then
							wgt_lon = min(rin(i+1,j,k), rout(i,j,k))
						else
							wgt_lon = min(rin(i,j,k), rout(i+1,j,k))
						end if
						qmfx(i,j,k) = wgt_lon * qmfx(i,j,k) + (1.0_r8 - wgt_lon) * fx_low(i,j,k)
					end do
				end do
			end do
			qmfx(:,mesh%full_lat_ibeg,:) = 0.0
			qmfx(:,mesh%full_lat_iend,:) = 0.0

			do k = mesh%full_lev_ibeg, mesh%full_lev_iend
				do j = mesh%half_lat_ibeg, mesh%half_lat_iend
					do i = mesh%full_lon_ibeg, mesh%full_lon_iend
						if (qmfy(i,j,k) - fy_low(i,j,k) >= 0.0) then
							wgt_lat = min(rin(i,j+1,k), rout(i,j,k))
						else
							wgt_lat = min(rin(i,j,k), rout(i,j+1,k))
						end if 
						qmfy(i,j,k) = wgt_lat * qmfy(i,j,k) + (1.0 - wgt_lat) * fy_low(i,j,k)
					end do
				end do
			end do
		end if
		call fill_zonal_halo_lon(qmfx, all_halo=.true.)
		end associate

	end subroutine calc_tracer_mass_flux_hori_upwind

	subroutine calc_tracer_mass_flux_vert_upwind(old_tracer, old_state, old_q, q, we, qmfz, dt, do_limiter)

		type(tracer_type), intent(inout) :: old_tracer
		type(state_type), intent(inout) :: old_state
		real(r8), intent(in) :: old_q(mesh%full_lon_lb:mesh%full_lon_ub, &
			 	 			  								  mesh%full_lat_lb:mesh%full_lat_ub, &
			 	 			  								  mesh%full_lev_lb:mesh%full_lev_ub)
		real(r8), intent(inout) :: q(mesh%full_lon_lb:mesh%full_lon_ub, &
			 	 			  								 mesh%full_lat_lb:mesh%full_lat_ub, &
			 	 			  								 mesh%full_lev_lb:mesh%full_lev_ub)
		real(r8), intent(in) :: we(mesh%full_lon_lb:mesh%full_lon_ub, &
															 mesh%full_lat_lb:mesh%full_lat_ub, &
															 mesh%half_lev_lb:mesh%half_lev_ub)
		real(r8), intent(out) :: qmfz(mesh%full_lon_lb:mesh%full_lon_ub, &
																  mesh%full_lat_lb:mesh%full_lat_ub, &
																  mesh%half_lev_lb:mesh%half_lev_ub)
		real(8), intent(in) :: dt
    logical, intent(in), optional :: do_limiter
		integer i, j, k
		real(r8) q_max, q_min, pin, qin, pout, qout, wgt_lev

		associate (fz_low => old_tracer%fz_low, &
							 q_td   => old_tracer%q_td  , &
							 rin    => old_tracer%rin   , &
							 rout   => old_tracer%rout  )

		do k = mesh%half_lev_ibeg + 1, mesh%half_lev_iend - 1
			do j = mesh%full_lat_ibeg, mesh%full_lat_iend
				do i = mesh%full_lon_ibeg, mesh%full_lon_iend
					if (k == mesh%half_lev_ibeg + 1 .or. k == mesh%half_lev_iend - 1) then
						qmfz(i,j,k) = we(i,j,k) * 0.5 * (q(i,j,k-1) + q(i,j,k)) * dt
				  else
					  qmfz(i,j,k) = we(i,j,k) * (7. / 12. * (q(i,j,k-1) + q(i,j,k)) - 1. / 12. * (q(i,j,k-2) + q(i,j,k+1)) + &
				  							sign(1._r8, we(i,j,k)) / 12. * (q(i,j,k+1) - q(i,j,k-2) - 3. * (q(i,j,k) - q(i,j,k-1)))) * dt
				  end if
				end do
			end do
		end do
		qmfz(:,:,mesh%half_lev_ibeg) = 0.0
		qmfz(:,:,mesh%half_lev_iend) = 0.0
		if (present(do_limiter) .and. do_limiter) then
			fz_low = 0.0
			do k = mesh%half_lev_ibeg + 1, mesh%half_lev_iend - 1
			  do j = mesh%full_lat_ibeg, mesh%full_lat_iend
					do i = mesh%full_lon_ibeg, mesh%full_lon_iend
						fz_low(i,j,k) = we(i,j,k) * (0.5 * (q(i,j,k-1) + q(i,j,k)) - sign(1._r8, we(i,j,k)) * 0.5 * (q(i,j,k) - q(i,j,k-1))) * dt
					end do
			  end do
		  end do
		  do k = mesh%full_lev_ibeg, mesh%full_lev_iend
		  	do j = mesh%full_lat_ibeg, mesh%full_lat_iend
		  		do i = mesh%full_lon_ibeg, mesh%full_lon_iend
		  			q_td(i,j,k) = (old_q(i,j,k) * old_state%m(i,j,k) - (fz_low(i,j,k+1) - fz_low(i,j,k))) / old_state%m(i,j,k)
		  		end do
		  	end do
		  end do

		  do k = mesh%full_lev_ibeg, mesh%full_lev_iend
		  	do j = mesh%full_lat_ibeg, mesh%full_lat_iend
		  		do i = mesh%full_lon_ibeg, mesh%full_lon_iend
		  			if (k == mesh%full_lev_ibeg) then
		  				q_max = max(maxval(old_q(i,j,k:k+1)), maxval(q_td(i,j,k:k+1)))
		  				q_min = min(minval(old_q(i,j,k:k+1)), minval(q_td(i,j,k:k+1)))
		  			else if (k == mesh%full_lev_iend) then
		  				q_max = max(maxval(old_q(i,j,k-1:k)), maxval(q_td(i,j,k-1:k)))
		  				q_min = min(minval(old_q(i,j,k-1:k)), minval(q_td(i,j,k-1:k)))
		  			else
		  				q_max = max(maxval(old_q(i,j,k-1:k+1)), maxval(q_td(i,j,k-1:k+1)))
		  				q_min = min(minval(old_q(i,j,k-1:k+1)), minval(q_td(i,j,k-1:k+1)))
		  			end if  
		  			pin   = max(0.0, qmfz(i,j,k  ) - fz_low(i,j,k  )) - min(0.0, qmfz(i,j,k+1) - fz_low(i,j,k+1))
		  			pout  = max(0.0, qmfz(i,j,k+1) - fz_low(i,j,k+1)) - min(0.0, qmfz(i,j,k  ) - fz_low(i,j,k  ))
		  			qin   = (q_max - q_td(i,j,k)) * old_state%m(i,j,k)
		  			qout  = (q_td(i,j,k) - q_min) * old_state%m(i,j,k)
		  			if (pin > 0) then
		  			  rin(i,j,k) = min(1.0, qin / pin)
		  			else
		  				rin(i,j,k) = 0.0
		  			end if
		  			if (pout > 0) then
		  				rout(i,j,k) = min(1.0, qout / pout)
		  			else
		  				rout(i,j,k) = 0.0
		  			end if
		  		end do
		  	end do
		  end do

		  do k = mesh%half_lev_ibeg + 1, mesh%half_lev_iend - 1
		  	do j = mesh%full_lat_ibeg, mesh%full_lat_iend
		  		do i = mesh%full_lon_ibeg, mesh%full_lon_iend
		  			if (qmfz(i,j,k) - fz_low(i,j,k) >= 0.0) then
		  				wgt_lev = min(rin(i,j,k), rout(i,j,k-1))
		  			else
		  				wgt_lev = min(rin(i,j,k-1), rout(i,j,k))
		  			end if
		  			qmfz(i,j,k) = (wgt_lev * qmfz(i,j,k) + (1.0 - wgt_lev) * fz_low(i,j,k))
		  		end do
		  	end do
		  end do
		end if
		end associate

	end subroutine calc_tracer_mass_flux_vert_upwind

	subroutine calc_hori_flux_upwind(tracer, state, q, dt, order, fx, fy)

		type(tracer_type), intent(in) :: tracer
		type(state_type), intent(in) :: state
		real(r8), intent(in) :: q(mesh%full_lon_lb:mesh%full_lon_ub, &
												      mesh%full_lat_lb:mesh%full_lat_ub, &
											 	      mesh%full_lev_lb:mesh%full_lev_ub)
		real(8), intent(in) :: dt
		character(len=*), intent(in) :: order
		real(r8), intent(out) :: fx(mesh%half_lon_lb:mesh%half_lon_ub, &
														    mesh%full_lat_lb:mesh%full_lat_ub, &
														    mesh%full_lev_lb:mesh%full_lev_ub)
		real(r8), intent(out) :: fy(mesh%full_lon_lb:mesh%full_lon_ub, &
														    mesh%half_lat_lb:mesh%half_lat_ub, &
														    mesh%full_lev_lb:mesh%full_lev_ub)
		integer i, j, k

		associate (u => tracer%u, v => tracer%v)
		select case (trim(order))
		case ('1st')
			do k = mesh%full_lev_ibeg, mesh%full_lev_iend
				do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
					do i = mesh%half_lon_ibeg, mesh%half_lon_iend
						fx(i,j,k) = u(i,j,k) * (0.5 * (q(i+1,j,k) + q(i,j,k))  - &
					  0.5 * sign(1._r8, u(i,j,k)) * (q(i+1,j,k) - q(i,j,k)))
					end do
				end do
			end do

			do k = mesh%full_lev_ibeg, mesh%full_lev_iend
				do j = mesh%half_lat_ibeg, mesh%half_lat_iend
					do i = mesh%full_lon_ibeg, mesh%full_lon_iend
						fy(i,j,k) = v(i,j,k) * (0.5 * (q(i,j+1,k) + q(i,j,k)) - &
						0.5 * sign(1._r8, v(i,j,k)) * (q(i,j+1,k) - q(i,j,k)))
					end do
				end do
			end do
	  case ('3rd')
			do k = mesh%full_lev_ibeg, mesh%full_lev_iend
				do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
					do i = mesh%half_lon_ibeg, mesh%half_lon_iend
						fx(i,j,k) = u(i,j,k) * (7.0 / 12.0 * (q(i+1,j,k) + q(i,j,k)) - 1.0 / 12.0 * (q(i+2,j,k) + q(i-1,j,k)) + &
					  			sign(1._r8, u(i,j,k)) / 12.0 * (q(i+2,j,k) - q(i-1,j,k) - 3.0 * (q(i+1,j,k) - q(i,j,k))))
					end do
				end do
			end do

			do k = mesh%full_lev_ibeg, mesh%full_lev_iend
				do j = mesh%half_lat_ibeg, mesh%half_lat_iend
					do i = mesh%full_lon_ibeg, mesh%full_lon_iend
						if (j == mesh%half_lat_ibeg .or. j == mesh%half_lat_iend) then
							fy(i,j,k) = v(i,j,k) * 0.5 * (q(i,j+1,k) + q(i,j,k)) ! - 0.5 * sign(1._r8, v(i,j,k)) * (mx(i,j+1,k) - mx(i,j,k)))
						else
						  fy(i,j,k) = v(i,j,k) * (7.0 / 12.0 * (q(i,j+1,k) + q(i,j,k)) - 1.0 / 12.0 * (q(i,j+2,k) + q(i,j-1,k)) + &
									  sign(1._r8, v(i,j,k)) / 12.0 * (q(i,j+2,k) - q(i,j-1,k) - 3.0 * (q(i,j+1,k) - q(i,j,k))))
						end if
					end do
				end do
			end do
		case default
			print*, "No such upwind scheme!"
			stop
		end select
		call fill_zonal_halo_lon(fx, all_halo=.true.)
		end associate

	end subroutine calc_hori_flux_upwind

end module upwind_mod